"""
Metric calculation and decomposition utilities.
"""
import pandas as pd
from typing import Dict, List, Optional, Any
from ..models.schemas import MetricConfig, DecompositionType, MetricDecompositionResult


class MetricCalculator:
    """Handles metric calculations and decompositions."""

    def __init__(self):
        self.metric_configs = {}

    def register_metric(self, config: MetricConfig):
        """Register a metric configuration."""
        self.metric_configs[config.name] = config

    def calculate_metric(
        self,
        df: pd.DataFrame,
        metric_name: str,
        config: Optional[MetricConfig] = None
    ) -> float:
        """
        Calculate a metric value from data.

        Args:
            df: DataFrame containing the data
            metric_name: Name of the metric to calculate
            config: Optional metric configuration

        Returns:
            Calculated metric value
        """
        config = config or self.metric_configs.get(metric_name)

        if config and config.formula:
            # If formula is defined, evaluate it
            return self._evaluate_formula(df, config.formula)
        elif metric_name in df.columns:
            # Direct column aggregation
            agg_func = config.aggregation if config else 'sum'
            if agg_func == 'sum':
                return df[metric_name].sum()
            elif agg_func == 'avg':
                return df[metric_name].mean()
            elif agg_func == 'max':
                return df[metric_name].max()
            elif agg_func == 'min':
                return df[metric_name].min()
            elif agg_func == 'count':
                return df[metric_name].count()
        else:
            raise ValueError(f"Metric {metric_name} not found in data or configuration")

    def _evaluate_formula(self, df: pd.DataFrame, formula: str) -> float:
        """
        Evaluate a metric formula.

        Simple formula evaluation for common patterns like:
        - revenue / users (ratio)
        - clicks / impressions (ratio)
        - revenue_a + revenue_b (sum)
        """
        # This is a simplified implementation
        # In production, you might want to use a safer expression evaluator
        try:
            # Create a namespace with column sums
            namespace = {col: df[col].sum() if col in df.columns else 0
                        for col in df.columns}
            return eval(formula, {"__builtins__": {}}, namespace)
        except Exception as e:
            raise ValueError(f"Failed to evaluate formula '{formula}': {str(e)}")

    def decompose_metric(
        self,
        current_df: pd.DataFrame,
        baseline_df: pd.DataFrame,
        metric_config: MetricConfig
    ) -> MetricDecompositionResult:
        """
        Decompose a metric into its components and analyze contribution.

        Args:
            current_df: Current period data
            baseline_df: Baseline period data
            metric_config: Configuration for the metric to decompose

        Returns:
            MetricDecompositionResult with decomposition analysis
        """
        if not metric_config.components:
            raise ValueError(f"Metric {metric_config.name} has no components defined")

        decomposition_type = metric_config.decomposition_type or DecompositionType.ADDITIVE

        # Calculate current and baseline values for each component
        current_components = {}
        baseline_components = {}

        for component in metric_config.components:
            current_components[component] = self.calculate_metric(
                current_df, component
            )
            baseline_components[component] = self.calculate_metric(
                baseline_df, component
            )

        # Calculate total metric values
        current_total = self._calculate_total(
            current_components, decomposition_type
        )
        baseline_total = self._calculate_total(
            baseline_components, decomposition_type
        )

        # Calculate changes
        component_changes = {
            comp: current_components[comp] - baseline_components[comp]
            for comp in metric_config.components
        }

        # Calculate contributions based on decomposition type
        if decomposition_type == DecompositionType.ADDITIVE:
            contributions = self._additive_contribution(
                component_changes, current_total - baseline_total
            )
        elif decomposition_type == DecompositionType.MULTIPLICATIVE:
            contributions = self._multiplicative_contribution(
                current_components, baseline_components, current_total, baseline_total
            )
        else:  # RATIO
            contributions = self._ratio_contribution(
                current_components, baseline_components, metric_config
            )

        return MetricDecompositionResult(
            metric=metric_config.name,
            formula=metric_config.formula or "",
            decomposition_type=decomposition_type,
            components=current_components,
            component_changes=component_changes,
            component_contributions=contributions
        )

    def _calculate_total(
        self,
        components: Dict[str, float],
        decomposition_type: DecompositionType
    ) -> float:
        """Calculate total based on decomposition type."""
        values = list(components.values())
        if decomposition_type == DecompositionType.ADDITIVE:
            return sum(values)
        elif decomposition_type == DecompositionType.MULTIPLICATIVE:
            result = 1.0
            for v in values:
                result *= v
            return result
        else:  # RATIO
            if len(values) >= 2:
                return values[0] / values[1] if values[1] != 0 else 0
            return values[0] if values else 0

    def _additive_contribution(
        self,
        component_changes: Dict[str, float],
        total_change: float
    ) -> Dict[str, float]:
        """Calculate contribution for additive decomposition."""
        if total_change == 0:
            return {comp: 0.0 for comp in component_changes}

        return {
            comp: change / total_change if total_change != 0 else 0
            for comp, change in component_changes.items()
        }

    def _multiplicative_contribution(
        self,
        current: Dict[str, float],
        baseline: Dict[str, float],
        current_total: float,
        baseline_total: float
    ) -> Dict[str, float]:
        """Calculate contribution for multiplicative decomposition."""
        total_change = current_total - baseline_total
        if total_change == 0:
            return {comp: 0.0 for comp in current}

        contributions = {}
        for comp in current:
            # Simplified multiplicative attribution
            change_rate = (current[comp] - baseline[comp]) / baseline[comp] if baseline[comp] != 0 else 0
            contributions[comp] = change_rate

        return contributions

    def _ratio_contribution(
        self,
        current: Dict[str, float],
        baseline: Dict[str, float],
        metric_config: MetricConfig
    ) -> Dict[str, float]:
        """Calculate contribution for ratio decomposition."""
        components = list(current.keys())
        if len(components) < 2:
            return {comp: 1.0 for comp in components}

        numerator, denominator = components[0], components[1]

        # Calculate current and baseline ratios
        current_ratio = current[numerator] / current[denominator] if current[denominator] != 0 else 0
        baseline_ratio = baseline[numerator] / baseline[denominator] if baseline[denominator] != 0 else 0

        # Impact of numerator and denominator changes
        numerator_change = current[numerator] - baseline[numerator]
        denominator_change = current[denominator] - baseline[denominator]

        ratio_change = current_ratio - baseline_ratio

        if ratio_change == 0:
            return {numerator: 0.0, denominator: 0.0}

        # Approximate contribution (simplified)
        numerator_contribution = (numerator_change / baseline[denominator]) / ratio_change if baseline[denominator] != 0 and ratio_change != 0 else 0
        denominator_contribution = -(baseline[numerator] * denominator_change / (baseline[denominator] ** 2)) / ratio_change if baseline[denominator] != 0 and ratio_change != 0 else 0

        return {
            numerator: numerator_contribution,
            denominator: denominator_contribution
        }
